from keras.models import Sequential
from keras.layers import Convolution2D, MaxPooling2D, Activation, Dropout, Flatten, Dense

def lenet_5(keep_prob):
    model = Sequential()
    # add layers into net
    # layer1-conv,input size 28*28, output size 32*28*28
    model.add(Convolution2D(
        filters=32,
        kernel_size=5,
        strides=1,
        padding='same',
        batch_input_shape=(None, 1, 28, 28),
        data_format='channels_first'
    ))
    model.add(Activation('relu'))

    # layer1-pooling, input size32*28*28, output size 32*14*14
    model.add(MaxPooling2D(
        pool_size=2,
        strides=2,
        padding='same',
        data_format='channels_first'
    ))
    model.add(Dropout(keep_prob))

    # layer2-conv, input size 32*14*14, output size 64*14*14
    model.add(Convolution2D(
        filters=64,
        kernel_size=5,
        strides=1,
        padding='same',
        data_format='channels_first'
    ))
    model.add(Activation('relu'))

    # layer2-pooling, input size 64*14*14, output size 64*7*7
    model.add(MaxPooling2D(
        pool_size=2,
        strides=2,
        padding='same',
        data_format='channels_first'
    ))
    model.add(Dropout(keep_prob))

    # layer3-dense, input size 64*7*7, output size 1*1024
    model.add(Flatten())
    model.add(Dense(1024))
    model.add(Activation('relu'))

    # layer4-dense, input size 1*1024, output size 1*10
    model.add(Dense(10))
    model.add(Activation('softmax'))

    return model
